use std::collections::HashSet;

type Coord = (i32, i32);

#[derive(Debug, PartialEq, Eq)]
struct Waterfall {
    coords: HashSet<Coord>,
    max_height: i32,
    bottom_height: i32,
}

impl Waterfall {
    fn fall_sand_part1(&mut self) -> bool {
        let mut sand_coord = (500, 0);
        loop {
            let (x, y) = sand_coord;

            if y == self.max_height {
                return false;
            }

            let down_coord = (x, y + 1);
            if self.coords.contains(&down_coord) {
                let lef_down_coord = (x - 1, y + 1);
                if self.coords.contains(&lef_down_coord) {
                    let right_down_coord = (x + 1, y + 1);
                    if self.coords.contains(&right_down_coord) {
                        // Stop fall！
                        self.coords.insert(sand_coord);
                        return true;
                    } else {
                        sand_coord = right_down_coord;
                    }
                } else {
                    sand_coord = lef_down_coord;
                }
            } else {
                sand_coord = down_coord;
            }
        }
    }

    fn fall_sand_part2(&mut self) -> bool {
        if self.coords.contains(&(500, 0)) {
            return false;
        }

        let mut sand_coord = (500, 0);
        loop {
            let (x, y) = sand_coord;
            if y + 1 == self.bottom_height {
                self.coords.insert(sand_coord);
                return true;
            }

            let down_coord = (x, y + 1);
            if self.coords.contains(&down_coord) {
                let lef_down_coord = (x - 1, y + 1);
                if self.coords.contains(&lef_down_coord) {
                    let right_down_coord = (x + 1, y + 1);
                    if self.coords.contains(&right_down_coord) {
                        // Stop fall！
                        self.coords.insert(sand_coord);
                        return true;
                    } else {
                        sand_coord = right_down_coord;
                    }
                } else {
                    sand_coord = lef_down_coord;
                }
            } else {
                sand_coord = down_coord;
            }
        }
    }
}

fn parse_input() -> Waterfall {
    let mut coords_set = HashSet::<Coord>::new();
    let input = include_str!("../input.txt");
    input
        .lines()
        .for_each(|line| {
            let coords = line
                .split(" -> ")
                .map(|token| {
                    // 536,90
                    let mut iter = token.split(',');
                    let x = iter.next().unwrap();
                    let y = iter.next().unwrap();
                    (x.parse().unwrap(), y.parse().unwrap())
                })
                .collect::<Vec<Coord>>();
            
            for (begin, end) in Iterator::zip(coords.iter(), coords.iter().skip(1)) {
                let (bx, by) = *begin;
                let (ex, ey) = *end;
                if bx == ex {
                    // Veritical
                    let x = bx;
                    let (max_y, min_y) = if by > ey { (by, ey) } else { (ey, by) };
                    for y in min_y..=max_y {
                        coords_set.insert((x, y));
                    } 
                } else if by == ey {
                    // Hor
                    let y = by;
                    let (max_x, min_x) = if bx > ex { (bx, ex) } else { (ex, bx) };
                    for x in min_x..=max_x {
                        coords_set.insert((x, y));
                    }
                } else {
                    panic!("")
                }
            }
        });
    
    let max_height = coords_set
        .iter()
        .max_by_key(|(_, y)| y)
        .map(|(_, y)| *y)
        .unwrap();

    Waterfall { coords: coords_set, max_height, bottom_height: max_height + 2 }
}

fn part1() {
    let mut waterfall = parse_input();
    let mut sand_size = 0;
    while waterfall.fall_sand_part1() {
        sand_size += 1;
    }
    println!("Part 1: {}", sand_size);
}

fn part2() {
    let mut waterfall = parse_input();

    let mut sand_size = 0;
    while waterfall.fall_sand_part2() {
        sand_size += 1;
    }
    println!("Part 1: {}", sand_size);
}

fn main() {
    part1();
    part2();
}
